#!/usr/bin/env python
"""
:mod:`pynchon.bio.na`
=====================

Provides support for nucleotides, codons, nucleic acid sequences and alignments.
"""

import numpy as np
from pynchon.core.symbol import Symbol0dS1, SymbolVector0dS1, SymbolArray0dS1, \
                                Symbol2dS1, SymbolVector2dS1, SymbolArray2dS1, \
                                Symbol3dS1, SymbolVector3dS1, SymbolArray3dS1
                                
from pynchon.bio.aa import AA, AASeq, AAArr, AAAln
from pynchon.bio.data.alphabet import \
                        IUPAC_DNA, IUPAC_DNA_ALPHABET, IUPAC_DNA_CODONS, \
                        IUPAC_RNA, IUPAC_RNA_ALPHABET, IUPAC_RNA_CODONS, \
                        IUPAC_NA, IUPAC_NA_ALPHABET, IUPAC_NA_CODONS, \
                        IUPAC_PROTEIN_ALPHABET, IUPAC_GAP, IUPAC_PROTEIN_STOP
                                      
from pynchon.bio.data.code import GC_1, CC_NA, CC_NA_ALPHABET


class Nt(Symbol0dS1):
    """
    Nucleotide Symbol class.
    """
    metadata = (
                ('alphabet', IUPAC_NA_ALPHABET),
                ('elements', IUPAC_NA_ALPHABET),
                ('cc', CC_NA_ALPHABET)
                )
    fval = np.array(IUPAC_GAP[0], dtype=Symbol0dS1.data_type)
    
    def complement(self, cc=None, inplace=False):
        cc = cc or self.metadata['cc']
        return self.recode(cc, inplace)
    

class NtSeq(SymbolVector0dS1, Nt):
    """
    Nucleotide Sequence (SymbolVector) class.
    """
    
    def nt(self, i=0):
        """
        Returns a view of the i-th nucleotide.
        """
        nt = self[i:i+1]
        nt.shape = ()
        return nt.view(type=Nt)        
        
    def reverse(self):
        """
        Returns a reversed copy of self. Note that this returns a new 
        ``NtSeq`` instance.
        """
        return NtSeq(self[::-1])
    
    def complement(self, cc=None):
        """
        Returns a sequence which is the complement of self. 
        """
        cc = cc or self.metadata['cc']
        return self.recode(cc, inplace =False)       
    
    def codonseq(self):
        return self.asarray().view(dtype=Codon.data_type, type=CodonSeq)
    

class NtAln(SymbolArray0dS1, NtSeq):
    """
    Nucleotide Alignment (SymbolArray) class.
    """


class Codon(Symbol3dS1):
    """
    Codon Symbol class. A ``Codon`` is an instance of ``Symbol3dS1`` i.e. it has
    rank-3 elements of single letter strings. The default alphabet of ``Codon``
    is ``IUPAC_DNA_ALPHABET`` and the default elements are ``IUPAC_DNA_CODONS``.
    The reason for this is because the standard genetic codes e.g. ``GC_1`` are
    4-letter DNA to 20-letter protein sequences. 
    """
    metadata = (
                ('alphabet', IUPAC_DNA_ALPHABET),
                ('elements', IUPAC_DNA_CODONS),
                ('gc', GC_1)
                 )
    fval = np.array((IUPAC_GAP[0],) * 3, dtype=Symbol3dS1.data_type)
    
    def translate(self, start=False, gc=None):
        """
        Returns the translation of self as a ``AA`` instance, given the genetic
        code "gc" and whether it is a start codon "start. 
        
        Arguments:
        
            - start (``bool``) If true will try to translate self using the 
              start codon to amino acid mapping. Note that most codons cannot
              be start codons and will return the IUPAC_GAP symbol.
            - gc (``dict``) a mapping between codons and amino acid symbols. 
              For each codon two symbols should be provided, the second is used
              if "start" is ``True``. E.g. ```{('A','T','G'):('X','-')}``
        """
        gc = gc or self.metadata['gc']
        aa_letter = gc[self.item()][bool(start)]
        return AA(aa_letter)


class CodonSeq(SymbolVector3dS1, Codon):
    """
    Codon Sequence class.
    """
    
    def codon(self, i=0):
        """
        Returns a view of the i-th codon.
        """
        codon = self[i:i+1]
        codon.shape = ()
        return codon.view(type=Codon)
    
    def ntseq(self):
        """
        Returns a view of self as a ``NtSeq`` array.
        """
        raw_array = self.asarray()
        return raw_array.view(dtype=Nt.data_type, type=NtSeq)  
          
    def translate(self, orf=False, gc=None):
        """
        """
        start = True
        aas = []
        for i in xrange(len(self)):
            codon = self.codon(i)
            aa = codon.translate(start, gc) if orf else codon.translate(gc=gc)
            if orf:
                if aa == IUPAC_GAP:
                    continue
                if aa == IUPAC_PROTEIN_STOP:
                    break
                start = False  
            aas.append(aa)
        if aas:
            return AASeq(aas)
        else:
            addon = 'Probably because no start codon has been found and ' + \
                    '"orf" is ``True``.' if orf else ''
            raise ValueError('This sequence cannot be translated. ' + addon)
                
                        
class CodonAln(SymbolArray3dS1, CodonSeq):
    """
    Codon Alignment (SymbolArray) class.
    """

